import argparse
import datetime
import lib
import logging

S = lib.set_conf('sample')

target_host = 'localhost'
proxies_map = {'direct': 7432, 'pgbouncer': 6432, 'pgagoral': 7433}

proxies = proxies_map.keys()
# this is active session.
degrees = [str(i) for i in [4, 8, 16, 32, 64, 128]]

workloads = ['select1.sql', 'lots_data.sql']
protocol = ['simple']


def run_workload(port, workload, degree, duration, protocol):
    cmd = f'PGPASSWORD="{S.PGPASSWORD}" pgbench -n -h{S.PGHOST} -M{protocol} -f {workload} -p{port} -c{degree} -T{duration} postgres'
    logging.warning(cmd)
    return lib.run_pgbench_cmd(cmd)


def parse_args():
    global workloads, proxies, degrees, target_host, protocol
    parser = argparse.ArgumentParser(
        description="xproxy perf testing, port is decided by proxy type")
    parser.add_argument('--workload', type=str,
                        help=f"parts of {','.join(workloads)}")
    parser.add_argument('--proxies', type=str,
                        default=','.join(proxies),
                        help=f"parts of {','.join(proxies)}")
    parser.add_argument('--degree', type=str,
                        default=','.join(degrees),
                        help=f"format: {','.join(degrees)}")
    parser.add_argument('--host', type=str,
                        help='the host for testing',
                        default=target_host)
    parser.add_argument('--protocol', type=str, default='simple',
                        help="simple,extended,prepared")
    args = parser.parse_args()

    if args.workload:
        workloads = args.workload.split(',')
    if args.proxies:
        proxies = args.proxies.split(',')
    if args.degree:
        degrees = args.degree.split(',')

    target_host = args.host
    protocol = args.protocol.split(',')


if __name__ == '__main__':
    project = 'xproxy'
    duration = '10'  # seconds

    parse_args()

    c = lib.get_pg_cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS autotest(id serial primary key,
    project text, start_time timestamp, end_time timestamp, data jsonb)''')

    for workload in workloads:
        for degree in degrees:
            for p in protocol:
                for proxy in proxies:
                    start_time = datetime.datetime.now()
                    res = {'proxy': proxy,
                           'workload': workload,
                           'degree': int(degree),
                           'protocol': p}
                    res.update(run_workload(
                        proxies_map[proxy], workload, degree, duration, p))
                    end_time = datetime.datetime.now()
                    logging.warning(res)
                    res['osinfo'] = lib.get_extra_osparam()
                    lib.pg_save_autotest(project, start_time, end_time, res)
